import numpy as np
import cv2
import os
import sys, imageio
from pathlib import Path
sys.path.append(str(Path(__file__).resolve().parents[3]))
import matplotlib.pyplot as plt
from mind3d.utils.eagermot_utils.kitti_util import *


cats = ['Pedestrian', 'Car', 'Cyclist']
cat_ids = {cat: i for i, cat in enumerate(cats)}
COLORS = [(255, 0, 255), (122, 122, 255), (255, 0, 0), (0,0,255), (0,255,0)]
color=['red', 'blue', 'black', 'green', 'yellow', 'purple', 'pink']


def draw_bbox(img, bbox_2d, color, id, thickness=1):
    if bbox_2d is None:
        return
    cv2.rectangle(img,
                  (int(bbox_2d[0]), int(bbox_2d[1])), (int(bbox_2d[2]), int(bbox_2d[3])),
                  color, thickness)
    ct = [(bbox_2d[0] + bbox_2d[2]) / 2, (bbox_2d[1] + bbox_2d[3]) / 2]
    txt = '{}'.format(id)
    cv2.putText(img, txt, (int(ct[0]), int(ct[1])), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, 
                color, thickness=1, lineType=cv2.LINE_AA)

def draw_point_cloud(ax, points, title, axes=[0, 1, 2], point_size=0.6, xlim3d=None, ylim3d=None, zlim3d=None):
    """
    Convenient method for drawing various point cloud projections as a part of frame statistics.
    """
    # 设置xyz三个轴的点云范围
    axes_limits = [
        [-20, 80], # X axis range
        [-20, 20], # Y axis range
        [-3, 5]    # Z axis range
    ]
    points*=3.5
    axes_str = ['X', 'Y', 'Z']
    # 禁止显示背后的网格
    ax.grid(False)
    # 创建散点图[1]:xyz数据集，[2]:点云的大小，[3]:点云的反射率数据,[4]:为灰度显示
    ax.scatter(*np.transpose(points[:, axes]), s=point_size, c=points[:, 3], cmap='gray')
    if len(axes) > 2:
        # 设置限制角度
        ax.set_xlim3d(*axes_limits[axes[0]])
        ax.set_ylim3d(*axes_limits[axes[1]])
        ax.set_zlim3d(*axes_limits[axes[2]])
        # 将背景颜色设置为RGBA格式，目前的参数以透明显示
        ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0))
        # 设置z轴标题
    else:
        # 2D限制角度，只有xy轴
        ax.set_xlim(*axes_limits[axes[0]])
        ax.set_ylim(*axes_limits[axes[1]])
    # User specified limits
    if xlim3d!=None:
        ax.set_xlim3d(xlim3d)
    if ylim3d!=None:
        ax.set_ylim3d(ylim3d)
    if zlim3d!=None:
        ax.set_zlim3d(zlim3d)


def compute_3d_box_cam2(h, w, l, x, y, z, yaw):
    # 计算旋转矩阵
    R = np.array([[np.cos(yaw), 0, np.sin(yaw)], [0, 1, 0], [-np.sin(yaw), 0, np.cos(yaw)]])
    # 8个顶点的xyz
    x_corners = [l/2,l/2,-l/2,-l/2,l/2,l/2,-l/2,-l/2]
    y_corners = [0,0,0,0,-h,-h,-h,-h] 
    z_corners = [w/2,-w/2,-w/2,w/2,w/2,-w/2,-w/2,w/2]
    # 旋转矩阵点乘(3，8)顶点矩阵
    corners_3d_cam2 = np.dot(R, np.vstack([x_corners,y_corners,z_corners]))
    # 加上location中心点，得出8个顶点旋转后的坐标
    corners_3d_cam2 += np.vstack([x,y,z])
    return corners_3d_cam2


def draw_kitti_2d(seq, data_path, pred_path, save_path, frames=20):
    # 2d vis
    pred_path = pred_path + '/{}.txt'.format(seq)
    pred_file2d = open(pred_path, 'r')
    i=0
    preds=[]
    images=[]
    last_id=-1
    for line in pred_file2d.readlines():
        tmp = line[:-1].split(' ')
        frame_id = int(tmp[0])
        track_id = int(tmp[1])
        cat_id = cat_ids[tmp[2]]
        bbox = [float(tmp[6]), float(tmp[7]), float(tmp[8]), float(tmp[9])]
        score = float(tmp[17])
        if last_id!=frame_id:
            if frame_id==frames:
                break
            if last_id != -1:
                images.append(img)
            file_path = '{}/{:06d}.png'.format(data_path+'image_02/'+seq+"", frame_id)
            img = cv2.imread(file_path)
        # preds.append([bbox, frame_id, track_id])
        draw_bbox(img,bbox,color = COLORS[track_id % 5], id=track_id)
        last_id=frame_id
    imageio.mimsave(save_path+'2d.gif', images, fps=3)


def draw_kitti_3d(seq, data_path, pred_path, save_path, frames=20):
    calib = Calibration(data_path+"calib/"+seq+".txt")
    connections = [
            [0, 1], [1, 2], [2, 3], [3, 0],  # Lower plane parallel to Z=0 plane
            [4, 5], [5, 6], [6, 7], [7, 4],  # Upper plane parallel to Z=0 plane
            [0, 4], [1, 5], [2, 6], [3, 7]  # Connections between upper and lower planes
        ]
    pred_3d=open(pred_path + '{}.txt'.format(seq))
    det_3d=[]
    last_id=-1
    imgs=[]
    for line in pred_3d.readlines():
        tmp = line[:-1].split(' ')
        frame_id = int(tmp[0])
        track_id = int(tmp[1])
        cat_id = cat_ids[tmp[2]]
        # h,w,l,x,y,z
        bbox = [float(tmp[-8]),float(tmp[-7]),float(tmp[-6]),float(tmp[-5]), float(tmp[-4]), float(tmp[-3]), float(tmp[-2])]
        score = float(tmp[-1])
        if last_id!=frame_id:
            # 绘制3D点云数据，创建一个大小为20*10的图形画板
            if last_id !=-1:
                plt.axis('off')
                fig.savefig(save_path+"0.png")
                imgs.append(cv2.cvtColor(cv2.imread(save_path+"0.png"), cv2.COLOR_BGR2RGB))
                plt.close('all')
            if frames==frame_id:
                break
            point_cloud = np.fromfile((os.path.join(data_path, ("velodyne/{}/%06d.bin"%frame_id).format(seq))), dtype=np.float32).reshape(-1,4)
            fig = plt.figure(figsize=(10, 10))
            # # 在画板中添加1*1的网格的第一个子图，为3D图像
            ax = fig.add_subplot(111, projection='3d')
            # fig, ax = plt.subplots(figsize=(20,10))
            # 改变绘制图像的视角，即相机的位置，elev为Z轴角度，azim为(x,y)角度
            ax.view_init(30, 190)
            # 在画板中画出点云显示数据，point_cloud[::x]x值越大，显示的点越稀疏
            draw_point_cloud(ax, point_cloud[::2], "velo_points") #axes = [0,1]
            
        # ax.add_patch(plt.Rectangle((100, 100), 200, 100, color="blue", fill=False, linewidth=1))
        corners_3d_cam2=compute_3d_box_cam2(bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5], bbox[6])
        corners_3d_velo = calib.project_rect_to_velo(corners_3d_cam2.T).T
        
        for connection in connections:
            ax.plot(*corners_3d_velo[:, connection], c=color[track_id % 7], lw=2) # c=COLORS[det[1] % 5]
        last_id=frame_id
    imageio.mimsave(save_path+'3d.gif', imgs, fps=3)


if __name__ == '__main__':
    pred_path_3d="/data0/HR_dataset/JIANG/EagerMOT/storage/workspace/kitti/training/result/tracking_det_0_0_seg_0.0_0.9_bbox_0.01_0.01_kf_dist_2d_full_[-3.5_-0.3]_0.3_a3_3_h1_2_2d_age_3_3_cleaning_0_3d/"
    pred_path_2d="/data0/HR_dataset/JIANG/EagerMOT/storage/workspace/kitti/training/result/tracking_det_0_0_seg_0.0_0.9_bbox_0.01_0.01_kf_dist_2d_full_[-3.5_-0.3]_0.3_a3_3_h1_2_2d_age_3_3_cleaning_0_2d_projected_3d"
    data_path="/data0/HR_dataset/KITTI_tracking/kitti/training/"
    save_path="utils/eagermot_utils/"
    frames=20
    # draw_kitti_3d(seq="0001", data_path=data_path, pred_path=pred_path_3d, save_path=save_path, frames=20)
    draw_kitti_2d(seq="0001", data_path=data_path, pred_path=pred_path_2d, save_path=save_path, frames=20)
    # print(point_cloud.shape)